{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cedde396-5500-4993-aa6b-8b9beba04e4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 数据预处理，进行标准化和数据增强\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((28, 28)),  # 将图像缩放到28x28\n",
    "    transforms.RandomHorizontalFlip(),  # 数据增强：随机水平翻转\n",
    "    transforms.RandomRotation(10),  # 数据增强：随机旋转\n",
    "    transforms.ToTensor(),  # 转换为Tensor\n",
    "    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # RGB图像的标准化\n",
    "])\n",
    "\n",
    "# 读取数据集，假设数据集目录中有训练集和验证集文件夹\n",
    "train_data = datasets.ImageFolder(root='D:/Anaconda3/train_data', transform=transform)\n",
    "test_data = datasets.ImageFolder(root='D:/Anaconda3/test_data', transform=transform)\n",
    "\n",
    "# 创建DataLoader\n",
    "train_loader = DataLoader(train_data, batch_size=64, shuffle=True)\n",
    "test_loader = DataLoader(test_data, batch_size=64, shuffle=False)\n",
    "\n",
    "# 定义CNN模型\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self, num_layers=3, num_kernels=64, num_classes=10):  # 增加卷积层和卷积核\n",
    "        super(CNN, self).__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.num_kernels = num_kernels\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "        # 创建卷积层\n",
    "        self.conv_layers = self._create_conv_layers()\n",
    "        \n",
    "        # 假设输入图像28x28，经过卷积和池化后，我们需要计算全连接层的输入大小\n",
    "        self.fc = nn.Linear(self._get_conv_output_size(), num_classes)\n",
    "        self.dropout = nn.Dropout(p=0.5)  # 增加Dropout层，避免过拟合\n",
    "\n",
    "    def _create_conv_layers(self):\n",
    "        layers = []\n",
    "        in_channels = 3  # 假设输入是RGB图像\n",
    "        for _ in range(self.num_layers):\n",
    "            layers.append(nn.Conv2d(in_channels, self.num_kernels, kernel_size=3, padding=1))\n",
    "            layers.append(nn.ReLU())\n",
    "            layers.append(nn.MaxPool2d(kernel_size=2, stride=2))  # 增加池化层\n",
    "            in_channels = self.num_kernels\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def _get_conv_output_size(self):\n",
    "        # 测试输入一个假图像（28x28）通过卷积层，得到输出的尺寸\n",
    "        x = torch.zeros(1, 3, 28, 28)  # 假设输入为 28x28 RGB 图像\n",
    "        x = self.conv_layers(x)\n",
    "        return x.numel()  # 返回展平后的元素数量\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layers(x)\n",
    "        x = x.view(x.size(0), -1)  # 展平卷积层的输出\n",
    "        x = self.fc(x)\n",
    "        x = self.dropout(x)  # 使用Dropout进行正则化\n",
    "        return x\n",
    "# 初始化模型、损失函数和优化器\n",
    "model = CNN(num_layers=3, num_kernels=64)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# 训练模型\n",
    "def train(model, train_loader, criterion, optimizer, epochs=5):\n",
    "    model.train()\n",
    "    epoch_losses = []\n",
    "    epoch_accuracies = []\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        running_loss = 0.0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for images, labels in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            # 统计训练过程中的准确率\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "            running_loss += loss.item()\n",
    "        \n",
    "        epoch_loss = running_loss / len(train_loader)\n",
    "        epoch_accuracy = correct / total\n",
    "        epoch_losses.append(epoch_loss)\n",
    "        epoch_accuracies.append(epoch_accuracy)\n",
    "        \n",
    "        print(f'Epoch {epoch+1}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.4f}')\n",
    "\n",
    "    return epoch_losses, epoch_accuracies\n",
    "\n",
    "# 测试模型\n",
    "def test(model, test_loader):\n",
    "    model.eval()  # 设置为评估模式\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():  # 不需要计算梯度\n",
    "        for images, labels in test_loader:\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    accuracy = correct / total\n",
    "    print(f'Test Accuracy: {accuracy:.4f}')\n",
    "    return accuracy\n",
    "\n",
    "def plot_learning_curve(train_losses, train_accuracies, test_accuracies):\n",
    "    plt.figure(figsize=(15, 5))\n",
    "\n",
    "    # 绘制训练损失曲线\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.plot(train_losses, label='Train Loss', color='blue')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('Training Loss')\n",
    "    plt.legend()\n",
    "\n",
    "    # 绘制训练准确率曲线\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.plot(train_accuracies, label='Train Accuracy', color='red')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title('Training Accuracy')\n",
    "    plt.legend()\n",
    "\n",
    "    # 绘制测试准确率曲线\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.plot(test_accuracies, label='Test Accuracy', color='green')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title('Test Accuracy')\n",
    "    plt.legend()\n",
    "\n",
    "\n",
    "\n",
    "# 执行训练并记录学习曲线\n",
    "train_losses, train_accuracies = train(model, train_loader, criterion, optimizer, epochs=5)\n",
    "test_accuracies = []\n",
    "\n",
    "# 测试模型每个epoch\n",
    "for epoch in range(5):\n",
    "    test_accuracy = test(model, test_loader)\n",
    "    test_accuracies.append(test_accuracy)\n",
    "\n",
    "# 绘制学习曲线\n",
    "plot_learning_curve(train_losses, train_accuracies, test_accuracies)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
